import os
import requests
import json
import dotenv
from langchain.agents import create_tool_calling_agent, AgentExecutor
from langchain_community.tools import GoogleSerperRun
from langchain_community.tools.openai_dalle_image_generation import OpenAIDALLEImageGenerationTool
from langchain_community.utilities import GoogleSerperAPIWrapper
from langchain_community.utilities.dalle_image_generator import DallEAPIWrapper
from langchain_core.prompts import ChatPromptTemplate
from pydantic import Field, BaseModel
from langchain_core.tools import tool, Tool
from langchain_openai import ChatOpenAI

dotenv.load_dotenv()

@tool
def bocha_web_search_tool(query: str, count: int = 8) -> str:
    """
    使用Bocha Web Search API进行联网搜索，返回搜索结果的字符串。

    参数:
    - query: 搜索关键词
    - count: 返回的搜索结果数量

    返回:
    - 搜索结果的字符串形式
    """
    url = 'https://api.bochaai.com/v1/web-search'
    headers = {
        'Authorization': f'Bearer {os.getenv("BOCHA_API_KEY")}',  # 请替换为你的API密钥
        'Content-Type': 'application/json'
    }
    data = {
        "query": query,
        "freshness": "noLimit",  # 搜索的时间范围，例如 "oneDay", "oneWeek", "oneMonth", "oneYear", "noLimit"
        "summary": True,  # 是否返回长文本摘要总结
        "count": count
    }

    response = requests.post(url, headers=headers, json=data)

    if response.status_code == 200:
        # 返回给大模型的格式化的搜索结果文本
        # 可以自己对博查的搜索结果进行自定义处理
        return str(response.json())
    else:
        raise Exception(f"API请求失败，状态码: {response.status_code}, 错误信息: {response.text}")

@tool
def baidu_qianfan_img_generate(query: str) -> str:
    """这是百度千帆的ai生成图片的工具。输出是图片的路径"""
    url = "https://qianfan.baidubce.com/v2/images/generations"
    payload = json.dumps({
        "prompt": query,
        "model": "irag-1.0"
    }, ensure_ascii=False)
    gaode_api_key = os.getenv("QIANFAN_API_KEY")
    headers = {
        'Content-Type': 'application/json',
        'Authorization': f'Bearer {gaode_api_key}'
    }
    response = requests.request("POST", url, headers=headers, data=payload.encode("utf-8"))
    if response.status_code == 200:
        # 返回给大模型的格式化的搜索结果文本
        # 可以自己对博查的搜索结果进行自定义处理
        return str(response.json())
    else:
        raise Exception(f"API请求失败，状态码: {response.status_code}, 错误信息: {response.text}")

class GoogleSerperArgsSchema(BaseModel):
    query: str = Field(description="执行谷歌搜索的查询语句")


class DallEArgsSchema(BaseModel):
    query: str = Field(description="输入应该是生成图像的文本提示(prompt)")

# 1. 创建博查查询工具
bocha_tool = Tool(
    name="BochaWebSearch",
    func=bocha_web_search_tool,
    description="使用Bocha Web Search API进行网络搜索",
    args_schema = GoogleSerperArgsSchema,
    api_wrapper = GoogleSerperAPIWrapper(),
)

# 2. 创建千帆文生图工具
qianfan_tool = Tool(
    name="QianfanImgGenerator",
    func=baidu_qianfan_img_generate,
    description="使用QIAN FAN irag-1.0生成图片",
    args_schema = DallEArgsSchema,
    # api_wrapper = GoogleSerperAPIWrapper(),
)

# 3.定义谷歌搜索工具
google_serper = GoogleSerperRun(
    name="google_serper",
    description=(
        "一个低成本的谷歌搜索API。"
        "当你需要回答有关时事的问题时，可以调用该工具。"
        "该工具的输入是搜索查询语句。"
    ),
    api_wrapper=GoogleSerperAPIWrapper(),
)

# 4. 定义OpenAI文生图工具
dalle = OpenAIDALLEImageGenerationTool(
    name="openai_dalle",
    api_wrapper=DallEAPIWrapper(model="dall-e-3"),
    # args_schema=DallEArgsSchema,
)
# tools = [google_serper, dalle]

tools = [bocha_tool, qianfan_tool]

# 2.定义工具调用agent提示词模板
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是由OpenAI开发的聊天机器人，善于帮助用户解决问题。"),
    ("placeholder", "{chat_history}"),
    ("human", "{input}"),
    ("placeholder", "{agent_scratchpad}"),
])

# 3.创建大语言模型  这里还是moonshot 如果用OpenAI -> 模型改成 ：gpt-4o-mini
llm = ChatOpenAI(model="kimi-k2-turbo-preview")

# 4.创建agent与agent执行者
agent = create_tool_calling_agent(
    prompt=prompt,
    llm=llm,
    tools=tools,
)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

print(agent_executor.invoke({"input": "帮我绘制一幅鲨鱼在天上游泳的场景"}))
